// Copyright 2022 Chen Jun

#include "armor_tracker/extended_kalman_filter.hpp"

namespace rm_auto_aim
{
  /*f是预测函数，h是观测函数，j_f是换衣函数的雅可比矩阵，j_h是观测函数的雅可比矩阵，u_q过程噪声，u_r测量噪声
  ，P0初始协方差矩阵，P_post后验协方差矩阵，n是状态向量的维度，I单位矩阵，x_pri是先验单位向量，x_post为后验状态向量*/
ExtendedKalmanFilter::ExtendedKalmanFilter(
  const VecVecFunc & f, const VecVecFunc & h, const VecMatFunc & j_f, const VecMatFunc & j_h,
  const VoidMatFunc & u_q, const VecMatFunc & u_r, const Eigen::MatrixXd & P0)
: f(f),
  h(h),
  jacobian_f(j_f),
  jacobian_h(j_h),
  update_Q(u_q),
  update_R(u_r),
  P_post(P0),
  n(P0.rows()),
  I(Eigen::MatrixXd::Identity(n, n)),
  x_pri(n),
  x_post(n)
{
}
//初始状态向量 x_post。
void ExtendedKalmanFilter::setState(const Eigen::VectorXd & x0) { x_post = x0; }
/*predict 函数用于执行预测步骤。它首先通过调用状态转换函数的雅可比矩阵来计算预测的状态转移矩阵 F，
然后通过调用过程噪声协方差函数来计算过程噪声协方差矩阵 Q。接下来，它使用状态转换函数 f 对当前状态 x_post 进行预测，
得到预测状态向量 x_pri。然后，通过使用预测的状态转移矩阵 F 和过程噪声协方差矩阵 Q 对先前估计的协方差矩阵 P_post 进行预测，
得到预测的协方差矩阵 P_pri。最后，它将预测状态向量 x_pri 和预测协方差矩阵 P_pri 分别赋值给当前状态向量 x_post 和
当前协方差矩阵 P_post。

*/
Eigen::MatrixXd ExtendedKalmanFilter::predict()
{
  F = jacobian_f(x_post), Q = update_Q();

  x_pri = f(x_post);
  P_pri = F * P_post * F.transpose() + Q;

  // handle the case when there will be no measurement before the next predict
  x_post = x_pri;
  P_post = P_pri;

  return x_pri;
}
/*update 函数用于执行更新步骤。它首先通过调用观测函数的雅可比矩阵来计算观测矩阵 H，然后通过调用测量噪声协方差函数
来计算测量噪声协方差矩阵 R。接下来，它使用预测状态向量 x_pri 和观测矩阵 H 来计算卡尔曼增益矩阵 K。然后，它使用观测
向量 z 和观测函数 h 来计算新的状态向量 x_post。最后，它通过使用卡尔曼增益矩阵 K 和观测矩阵 H 对先前估计的协方差矩阵P_pri
 进行更新，得到新的协方差矩阵 P_post。*/
Eigen::MatrixXd ExtendedKalmanFilter::update(const Eigen::VectorXd & z)
{
  H = jacobian_h(x_pri), R = update_R(z);

  K = P_pri * H.transpose() * (H * P_pri * H.transpose() + R).inverse();
  x_post = x_pri + K * (z - h(x_pri));
  P_post = (I - K * H) * P_pri;

  return x_post;
}

}  // namespace rm_auto_aim
